package cryptotools

import (
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"io/ioutil"
	"os"
	"strings"
)

type KeyPair struct {
	PrivateKey rsa.PrivateKey
	PublicKey  rsa.PublicKey
}

func RSA_GenKey(bits int) KeyPair {
	privateKey, _ := rsa.GenerateKey(rand.Reader, bits)
	publicKey := privateKey.PublicKey
	return KeyPair{PrivateKey: *privateKey, PublicKey: publicKey}
}

func RSA_SavePublicKey(saveDir string, filenameWithSuffix string, publicKey *rsa.PublicKey) {
	saveDir, _ = strings.CutSuffix(saveDir, "/")
	publicKeyFile, err := os.OpenFile(saveDir+"/"+filenameWithSuffix, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600)
	if err != nil {
		panic(err)
	}
	defer publicKeyFile.Close()
	publicKeyPEM := &pem.Block{
		Type:  "RSA PUBLIC KEY",
		Bytes: x509.MarshalPKCS1PublicKey(publicKey),
	}
	err = pem.Encode(publicKeyFile, publicKeyPEM)
	if err != nil {
		panic(err)
	}
}

func RSA_SavePrivateKey(saveDir string, filenameWithSuffix string, privateKey *rsa.PrivateKey) {
	saveDir, _ = strings.CutSuffix(saveDir, "/")
	privateKeyFile, err := os.OpenFile(saveDir+"/"+filenameWithSuffix, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0600)
	if err != nil {
		panic(err)
	}
	defer privateKeyFile.Close()
	privateKeyPEM := &pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(privateKey),
	}
	err = pem.Encode(privateKeyFile, privateKeyPEM)
	if err != nil {
		panic(err)
	}
}

func RSA_LoadPublicKey(saveDir string, fileNameWithSuffix string) *rsa.PublicKey {
	saveDir, _ = strings.CutSuffix(saveDir, "/")
	data, err := os.ReadFile(saveDir + "/" + fileNameWithSuffix)
	if err != nil {
		panic(err)
	}
	block, _ := pem.Decode(data)
	if block == nil {
		fmt.Errorf("failed to parse PEM block")
		return nil
	}
	if block.Type != "RSA PUBLIC KEY" {
		fmt.Errorf("Not a public key file")
		return nil
	}
	result, _ := x509.ParsePKCS1PublicKey(block.Bytes)
	return result
}

func RSA_LoadPrivateKey(saveDir string, fileNameWithSuffix string) *rsa.PrivateKey {
	saveDir, _ = strings.CutSuffix(saveDir, "/")
	data, err := ioutil.ReadFile(saveDir + "/" + fileNameWithSuffix)
	if err != nil {
		panic(err)
	}
	block, _ := pem.Decode(data)
	if block == nil {
		fmt.Errorf("failed to parse PEM block")
		return nil
	}
	if block.Type != "RSA PRIVATE KEY" {
		fmt.Errorf("Not a private key file")
		return nil
	}
	result, _ := x509.ParsePKCS1PrivateKey(block.Bytes)
	return result
}

func RSA_Encrypt(plainText []byte, publicKey *rsa.PublicKey) ([]byte, error) {
	cipherText, err := rsa.EncryptPKCS1v15(rand.Reader, publicKey, plainText)
	if err != nil {
		return []byte{}, err
	}
	return cipherText, nil
}

func RSA_Decrypt(cipherText []byte, privateKey *rsa.PrivateKey) ([]byte, error) {
	plainText, err := rsa.DecryptPKCS1v15(rand.Reader, privateKey, cipherText)
	if err != nil {
		return []byte{}, err
	}
	return plainText, nil
}

func RSA_Sign(text []byte, privateKey *rsa.PrivateKey) ([32]byte, []byte) {
	hashed := sha256.Sum256(text)
	signature, err := rsa.SignPKCS1v15(rand.Reader, privateKey, crypto.SHA256, hashed[:])
	if err != nil {
		panic(err)
	}
	return hashed, signature
}

func RSA_SignVerify(hashed [32]byte, signature []byte, publicKey *rsa.PublicKey) bool {
	err := rsa.VerifyPKCS1v15(publicKey, crypto.SHA256, hashed[:], signature)
	return (err == nil)
}
